import os
import shutil
import random

# 设置目录路径
data_dir = './datasets'
images_dir = os.path.join(data_dir, 'images')
labels_dir = os.path.join(data_dir, 'labels')

# 确保目标目录存在
def ensure_dir(dir_path):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

# 创建训练、验证、测试子目录
for split in ['train', 'val', 'test']:
    ensure_dir(os.path.join(images_dir, split))
    ensure_dir(os.path.join(labels_dir, split))

# 获取所有图像文件名
image_files = [f for f in os.listdir(images_dir) if f.endswith('.png')]

# 设置划分比例
train_ratio = 0.9
val_ratio = 0.05
test_ratio = 0.05

# 打乱文件顺序
random.seed(42)  # 保证结果可复现
random.shuffle(image_files)

# 按比例划分
train_count = int(len(image_files) * train_ratio)
val_count = int(len(image_files) * val_ratio)

train_files = image_files[:train_count]
val_files = image_files[train_count:train_count + val_count]
test_files = image_files[train_count + val_count:]

# 移动文件到相应的子目录
def move_files(files, split):
    for file in files:
        image_path = os.path.join(images_dir, file)
        label_file = file.replace('.png', '.txt')
        label_path = os.path.join(labels_dir, label_file)
        
        # 确保标签文件存在才移动
        if os.path.exists(label_path):
            shutil.move(image_path, os.path.join(images_dir, split, file))
            shutil.move(label_path, os.path.join(labels_dir, split, label_file))
        else:
            print(f"标签文件 {label_file} 不存在，跳过该图片 {file}")

# 分别将训练、验证、测试文件移动到对应目录
move_files(train_files, 'train')
move_files(val_files, 'val')
move_files(test_files, 'test')

print("数据集划分完成！")
